# routes/knowledge.py
from flask import Blueprint, request, jsonify
from app.utils.knowledge_base import (
    build_vectorstore_from_text,
    save_vectorstore_to_disk,
    load_vectorstore_from_disk
)
import os

knowledge_bp = Blueprint("knowledge", __name__, url_prefix="/knowledge")

VECTORSTORE_DIR = "vectorstores"
os.makedirs(VECTORSTORE_DIR, exist_ok=True)

@knowledge_bp.route("/upload", methods=["POST"])
def upload_knowledge():
    if "file" not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files["file"]
    filename = file.filename.strip()

    if not filename.endswith(".txt"):
        return jsonify({"error": "Only .txt files are supported"}), 400

    try:
        text = file.read().decode("utf-8")
        vectorstore = build_vectorstore_from_text(text)

        save_name = filename.rsplit(".", 1)[0]
        save_path = os.path.join(VECTORSTORE_DIR, save_name)
        save_vectorstore_to_disk(vectorstore, save_path)

        return jsonify({"message": "知识库构建成功", "vector_id": save_name}), 200

    except Exception as e:
        return jsonify({"error": str(e)}), 500


@knowledge_bp.route("/query", methods=["POST"])
def query_knowledge():
    data = request.get_json()
    question = data.get("question", "").strip()
    vector_id = data.get("vector_id", "").strip()

    if not question or not vector_id:
        return jsonify({"error": "Missing question or vector_id"}), 400

    try:
        path = os.path.join(VECTORSTORE_DIR, vector_id)
        if not os.path.isdir(path):
            return jsonify({"error": "知识库不存在"}), 404

        vectorstore = load_vectorstore_from_disk(path)
        retriever = vectorstore.as_retriever()
        docs = retriever.invoke(question)

        context = "\n\n".join([doc.page_content for doc in docs])

        return jsonify({
            "question": question,
            "context": context,
            "status": "success"
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500